import numpy as np
import numpy.fft as fft
import matplotlib.pyplot as plt
import os

class WaveEquationSolver:
    def __init__(self, Lx, dx, total_t, a, c,ndiag):
        self.Lx = Lx
        self.dx = dx
        self.Nx = int(Lx / dx) + 1
        self.x = np.linspace(0, Lx, self.Nx)
        self.dt = c * self.dx / a
        self.total_t = total_t
        self.ntime = int(total_t / self.dt) + 1
        self.a = a
        self.c = c
        self.ndiag=ndiag
        self.u = np.zeros_like(self.x)
        self.set_initial_conditions()
        self.out_iter=0

    def set_initial_conditions(self):
        self.u = np.sin(2 * np.pi * self.x)
        self.u[0] = self.u[-2]
        self.u[-1] = self.u[1]

    def forward_u_LW1(self):
        # 使用Lax-Wendroff格式进行一步时间步进
        self.u[1:-1] = self.u[1:-1] - 0.5 * self.c * (self.u[2:] - self.u[:-2]) + 0.5 * self.c**2 * (self.u[2:] - 2 * self.u[1:-1] + self.u[:-2])

    def set_period_bnd(self):
        self.u[0] = self.u[-2]
        self.u[-1] = self.u[1]

    def run(self,flag_plot, flag_spectrum):
        self.out_iter = 0
        while self.out_iter < self.ntime:
            self.out_iter += 1
            self.forward_u_LW1()
            self.set_period_bnd()
            if self.out_iter % self.ndiag == 0:
                if flag_plot:
                    self.plot_solution(flag_show=False,flag_save=False)
                if flag_spectrum:
                    self.plot_spectrum()

    def plot_solution(self,flag_show,flag_save):
        # 绘制解的空间分布
        t_diag = self.out_iter * self.dt
        plt.figure()
        plt.plot(self.x, self.u)
        plt.xlabel('x')
        plt.ylabel('u')
        plt.title(f't={t_diag:.2f}, c={self.c:.2f}')
        plt.grid(True)
        if flag_show:
            plt.show()
        if flag_save:
            plt.savefig(f'u_x_plot_step_{self.out_iter}.png')
        plt.close()

    def plot_spectrum(self):
        # 绘制解的频谱
        t_diag = self.out_iter * self.dt
        u_spectrum = fft.fft(self.u)
        real_spectrum = np.abs(u_spectrum)[:int(self.Nx / 2 + 1)]
        freqs = fft.fftfreq(self.Nx, self.dx)
        non_negative_freqs = freqs[:int(self.Nx / 2 + 1)]
        plt.figure()
        plt.plot(non_negative_freqs, real_spectrum)
        plt.xlabel('Frequency')
        plt.ylabel('Amplitude')
        plt.title(f't={t_diag:.2f}, c={self.c:.2f}')
        plt.grid(True)
        plt.savefig(f'u_spectrum_plot_step_{self.out_iter}.png')
        plt.close()

# 使用示例
Lx = 3
dx = 0.02
total_t = 10
a = 1
c = 0.99
ndiag=10

solver = WaveEquationSolver(Lx, dx, total_t, a, c,ndiag)
solver.run(flag_plot=False, flag_spectrum=False)
solver.plot_solution(flag_show=True,flag_save=False)